#include <iostream>
#include <vector>
#include <string>
#include <filesystem>
#include <opencv2/opencv.hpp>
#include <sstream>

namespace fs = std::filesystem;

const std::string TARGET_FOLDER = "Target";
const std::string TARGET_SUFFIX = ".png";
const std::string TARGET_SUFFIX2 = ".PNG";
const std::string OUTPUT_PATH = "output.png";

std::vector<std::string>* GetAllPngInTarget()
{
	std::vector<std::string>* ret = new std::vector<std::string>();

	fs::path currentPath = fs::current_path() / TARGET_FOLDER;
	std::cout << "Target directory: " << currentPath << std::endl;

	for (const auto& entry : fs::directory_iterator(currentPath))
	{
		if (entry.is_regular_file() && (entry.path().extension() == TARGET_SUFFIX || entry.path().extension() == TARGET_SUFFIX2))
		{
			std::string str = (currentPath / entry.path().filename()).string();
			std::cout << "Found PNG file: " << str << std::endl;
			ret->push_back(str);
		}
	}

	return ret;
}

int main(int argc, char** argv)
{
	int row = 6;
	if (argc > 1)
	{
		std::cout << "Param num: " << argc << std::endl;

		for (int i = 0; i < argc; ++i)
		{
			std::cout << "Param " << i << ": " << argv[i] << std::endl;
		}

		std::stringstream ss(argv[1]);
		int num;
		ss >> num;
		std::cout << "Read param num: " << num << std::endl;
		if (num > 0)
		{
			row = num;
		}
	}

	std::cout << "Row num: " << row << std::endl;

	std::vector<std::string>* list = GetAllPngInTarget();

	int totalCnt = 0;

	cv::Mat mat;
	std::vector<cv::Mat> rowList;
	for (int i = 0; i < list->size(); i++)
	{
		cv::String path = (*list)[i];
		cv::Mat m = cv::imread(path);
		if (m.empty())
		{
			std::cerr << "[Err] Empty image: " << path << std::endl;
			return 0;
		}

		if (i % row == 0)
		{
			mat = m;
		}
		else
		{
			cv::hconcat(mat, m, mat);
		}

		if (i % row == row - 1)
		{
			rowList.push_back(mat);
		}
		else if (i == list->size() - 1)
		{
			cv::Mat empty = cv::Mat::ones(m.size().height, m.size().width, m.type());
			for (int j = i % row + 1; j < row; j++)
			{
				cv::hconcat(mat, empty, mat);
			}
			rowList.push_back(mat);
		}

		totalCnt++;
	}

	if (rowList.size() > 0)
	{
		cv::Mat output = rowList[0];
		for (int i = 1; i < rowList.size(); i++)
		{
			cv::vconcat(output, rowList[i], output);
		}

		cv::imwrite(OUTPUT_PATH, output);
	}

	std::cout << "Finished! Count: " << totalCnt << std::endl;
	std::cin.get();

	delete list;

	return 0;
}